Upload 24 files

__init__ (1).py

@@ -0,0 +1,26 @@
+"""spark-eval: an execution-based benchmark for PySpark code generation."""
+
+from .harness import ExecResult, compare_frames, evaluate_candidate, run_code
+from .prompting import build_prompt, extract_code
+from .runner import RunReport, format_report, get_spark, pass_at_k, run_suite
+from .schema import Task, TaskValidationError, load_task, load_tasks
+
+__version__ = "0.1.0"
+
+__all__ = [
+    "ExecResult",
+    "RunReport",
+    "Task",
+    "TaskValidationError",
+    "build_prompt",
+    "compare_frames",
+    "evaluate_candidate",
+    "extract_code",
+    "format_report",
+    "get_spark",
+    "load_task",
+    "load_tasks",
+    "pass_at_k",
+    "run_code",
+    "run_suite",
+]

agg_count_null_semantics.yaml

@@ -0,0 +1,42 @@
+id: agg_count_null_semantics
+category: aggregations
+difficulty: medium
+probes: >
+  count(*) counts rows, count(col) skips NULLs, countDistinct(col) skips NULLs
+  and dedups. Three counts that differ only because of NULL handling. Models
+  reach for count("*") for all three, or use count("col") where the prompt
+  asked for rows.
+tags: [count, null_semantics, distinct]
+
+prompt: |
+  For each dept in `staff`, return one row with:
+    - n_rows      the number of rows in the group
+    - n_emails    the number of rows whose email is not null
+    - n_distinct  the number of distinct non-null emails
+
+  Return columns: dept, n_rows, n_emails, n_distinct.
+
+fixtures:
+  - name: staff
+    schema: dept STRING, name STRING, email STRING
+    rows:
+      - ["eng", "ann", "a@x.com"]
+      - ["eng", "bob", null]
+      - ["eng", "cal", "a@x.com"]
+      - ["eng", "dee", "d@x.com"]
+      - ["ops", "eve", null]
+      - ["ops", "fay", null]
+
+solution: |
+  from pyspark.sql import functions as F
+
+  def solve(spark, staff):
+      return staff.groupBy("dept").agg(
+          F.count(F.lit(1)).alias("n_rows"),
+          F.count("email").alias("n_emails"),
+          F.countDistinct("email").alias("n_distinct"),
+      )
+
+compare:
+  mode: rows
+  check_schema: false

agg_pivot_fill.yaml

@@ -0,0 +1,44 @@
+id: agg_pivot_fill
+category: aggregations
+difficulty: medium
+probes: >
+  pivot produces NULL for absent combinations, and the pivot column values
+  determine the output schema. Pinning the value list is also the difference
+  between one Spark job and two -- an unpinned pivot triggers an extra scan to
+  discover distinct values.
+tags: [pivot, groupby, null_fill]
+
+prompt: |
+  Pivot `sales_long` so there is one row per store and one column per quarter,
+  with columns named exactly Q1, Q2, Q3, containing the sum of amount.
+  Missing combinations must be 0, not null. All three quarter columns must be
+  present even if a quarter never appears in the data.
+
+  Return columns: store, Q1, Q2, Q3.
+
+fixtures:
+  - name: sales_long
+    schema: store STRING, quarter STRING, amount INT
+    rows:
+      - ["s1", "Q1", 10]
+      - ["s1", "Q1", 5]
+      - ["s1", "Q2", 20]
+      - ["s2", "Q3", 7]
+
+solution: |
+  from pyspark.sql import functions as F
+
+  def solve(spark, sales_long):
+      # Passing the value list explicitly keeps Q3 in the schema for s1 and
+      # avoids the discovery pass over the data.
+      return (
+          sales_long
+          .groupBy("store")
+          .pivot("quarter", ["Q1", "Q2", "Q3"])
+          .agg(F.sum("amount"))
+          .fillna(0, subset=["Q1", "Q2", "Q3"])
+      )
+
+compare:
+  mode: rows
+  check_schema: false

cli.py

@@ -0,0 +1,151 @@
+"""Command line interface."""
+
+from __future__ import annotations
+
+import argparse
+import sys
+from pathlib import Path
+
+from .runner import format_report, get_spark, run_suite, write_report
+from .schema import TaskValidationError, load_tasks
+
+DEFAULT_TASKS = Path(__file__).resolve().parent.parent / "tasks"
+
+
+def _add_selection_args(p: argparse.ArgumentParser) -> None:
+    p.add_argument("--tasks", type=Path, default=DEFAULT_TASKS, help="task directory")
+    p.add_argument("--category", action="append", dest="categories", help="filter (repeatable)")
+    p.add_argument("--id", action="append", dest="ids", help="run specific task ids")
+
+
+def cmd_run(args: argparse.Namespace) -> int:
+    from .models import build_model  # local import: keeps `validate` dependency-light
+
+    tasks = load_tasks(args.tasks, args.categories, args.ids)
+    if not tasks:
+        print("no tasks matched", file=sys.stderr)
+        return 1
+
+    model = build_model(args.model, timeout=args.request_timeout)
+    ks = tuple(sorted({1, *(args.k or [])}))
+    if args.n < max(ks):
+        print(
+            f"error: --n {args.n} is too small for pass@{max(ks)}; "
+            f"the estimator needs n >= k",
+            file=sys.stderr,
+        )
+        return 2
+
+    report = run_suite(
+        tasks,
+        model,
+        n_samples=args.n,
+        ks=ks,
+        temperature=args.temperature,
+        max_tokens=args.max_tokens,
+        timeout=args.timeout,
+        keep_responses=not args.no_responses,
+    )
+
+    print(format_report(report))
+    if args.out:
+        write_report(report, args.out)
+        print(f"wrote {args.out}")
+    return 0
+
+
+def cmd_validate(args: argparse.Namespace) -> int:
+    """Structural checks only. `selfcheck` is the one that executes anything."""
+    try:
+        tasks = load_tasks(args.tasks, args.categories, args.ids)
+    except TaskValidationError as exc:
+        print(f"INVALID: {exc}", file=sys.stderr)
+        return 1
+    from collections import Counter
+
+    counts = Counter(t.category for t in tasks)
+    print(f"{len(tasks)} tasks, all structurally valid\n")
+    for cat, n in sorted(counts.items()):
+        print(f"  {cat:<22} {n:>4}")
+    missing = [t.id for t in tasks if not t.probes]
+    if missing:
+        print(f"\nwarning: {len(missing)} tasks have no 'probes' note: {missing[:5]}")
+    return 0
+
+
+def cmd_selfcheck(args: argparse.Namespace) -> int:
+    """Execute every reference solution.
+
+    This is the check that matters. If a gold solution does not run, every
+    model scored against that task gets a meaningless result.
+    """
+    from .harness import evaluate_candidate
+
+    tasks = load_tasks(args.tasks, args.categories, args.ids)
+    spark = get_spark("spark-eval-selfcheck")
+    spark.sparkContext.setLogLevel("ERROR")
+
+    failures = []
+    for i, task in enumerate(tasks, 1):
+        result = evaluate_candidate(spark, task, task.solution, timeout=args.timeout)
+        status = "ok" if result.ok else f"BROKEN ({result.status})"
+        print(f"[{i:>3}/{len(tasks)}] {task.id:<40} {status}", flush=True)
+        if not result.ok:
+            failures.append((task.id, result.detail))
+
+    spark.stop()
+
+    if failures:
+        print(f"\n{len(failures)} reference solution(s) failed:\n", file=sys.stderr)
+        for tid, detail in failures:
+            print(f"  {tid}: {detail}", file=sys.stderr)
+        return 1
+    print(f"\nall {len(tasks)} reference solutions execute and self-compare")
+    return 0
+
+
+def main(argv: list[str] | None = None) -> int:
+    parser = argparse.ArgumentParser(
+        prog="spark-eval",
+        description="Execution-based benchmark for PySpark code generation.",
+    )
+    sub = parser.add_subparsers(dest="command", required=True)
+
+    p_run = sub.add_parser("run", help="score a model against the suite")
+    _add_selection_args(p_run)
+    p_run.add_argument(
+        "--model",
+        required=True,
+        help="ollama:<tag> | openai:<model> | dummy:reference",
+    )
+    p_run.add_argument("--n", type=int, default=1, help="samples per task")
+    p_run.add_argument(
+        "--k", type=int, action="append", help="report pass@k (repeatable, needs n>=k)"
+    )
+    p_run.add_argument("--temperature", type=float, default=0.2)
+    p_run.add_argument("--max-tokens", type=int, default=1024)
+    p_run.add_argument("--timeout", type=int, default=60, help="per-task exec seconds")
+    p_run.add_argument("--request-timeout", type=int, default=300)
+    p_run.add_argument("--out", type=Path, help="write full JSON report here")
+    p_run.add_argument(
+        "--no-responses",
+        action="store_true",
+        help="omit raw generations from the report (smaller files)",
+    )
+    p_run.set_defaults(func=cmd_run)
+
+    p_val = sub.add_parser("validate", help="structural check on task files")
+    _add_selection_args(p_val)
+    p_val.set_defaults(func=cmd_validate)
+
+    p_self = sub.add_parser("selfcheck", help="execute every reference solution")
+    _add_selection_args(p_self)
+    p_self.add_argument("--timeout", type=int, default=60)
+    p_self.set_defaults(func=cmd_selfcheck)
+
+    args = parser.parse_args(argv)
+    return args.func(args)
+
+
+if __name__ == "__main__":
+    raise SystemExit(main())

conftest.py

@@ -0,0 +1,18 @@
+from __future__ import annotations
+
+from pathlib import Path
+
+import pytest
+
+from spark_eval.runner import get_spark
+
+TASKS_DIR = Path(__file__).resolve().parent.parent / "tasks"
+
+
+@pytest.fixture(scope="session")
+def spark():
+    """One session for the whole test run. JVM startup dominates otherwise."""
+    session = get_spark("spark-eval-tests")
+    session.sparkContext.setLogLevel("ERROR")
+    yield session
+    session.stop()

harness.py

@@ -0,0 +1,284 @@
+"""Execution harness: run candidate code against fixtures and judge the result.
+
+This is the part that makes spark-eval different from string-similarity scoring.
+Nothing here looks at the *text* of the generated code. It runs it and compares
+the DataFrame that comes out.
+
+Security note, stated plainly: `run_code` executes untrusted model output in
+this process. The import guard below stops casual accidents (a model that
+decides to `import os` and clean up after itself), not a determined adversary.
+If you are scoring untrusted checkpoints, run the CLI inside a container with no
+network and a read-only mount.
+"""
+
+from __future__ import annotations
+
+import builtins
+import math
+import signal
+from collections.abc import Iterator
+from contextlib import contextmanager
+from dataclasses import dataclass
+from typing import Any
+
+from pyspark.sql import DataFrame, SparkSession
+
+from .schema import Compare, Fixture, Task
+
+# Modules a correct PySpark answer never needs. Blocking them turns "the model
+# wandered off and deleted the fixtures" into a clean task failure.
+BLOCKED_IMPORTS = {
+    "os",
+    "sys",
+    "subprocess",
+    "shutil",
+    "socket",
+    "requests",
+    "urllib",
+    "urllib2",
+    "urllib3",
+    "httpx",
+    "aiohttp",
+    "pathlib",
+    "ctypes",
+    "importlib",
+    "pickle",
+    "multiprocessing",
+    "tempfile",
+    "glob",
+}
+
+
+class TaskTimeout(Exception):
+    pass
+
+
+class BlockedImport(Exception):
+    pass
+
+
+@dataclass
+class ExecResult:
+    """Outcome of running one piece of code against one task."""
+
+    ok: bool
+    # "pass" | "error" | "timeout" | "no_solve" | "wrong_type"
+    # | "schema_mismatch" | "row_mismatch" | "blocked_import"
+    status: str
+    detail: str = ""
+
+    def __bool__(self) -> bool:  # pragma: no cover - convenience only
+        return self.ok
+
+
+@contextmanager
+def _time_limit(seconds: int) -> Iterator[None]:
+    """Wall-clock cap on a block of code.
+
+    SIGALRM only interrupts the driver thread, so a candidate that wedges deep
+    inside a JVM call can outlive this. Per-task subprocess isolation is on the
+    roadmap; until then, treat the timeout as best-effort.
+    """
+
+    def _handler(signum, frame):  # noqa: ANN001
+        raise TaskTimeout(f"exceeded {seconds}s")
+
+    previous = signal.signal(signal.SIGALRM, _handler)
+    signal.alarm(seconds)
+    try:
+        yield
+    finally:
+        signal.alarm(0)
+        signal.signal(signal.SIGALRM, previous)
+
+
+def _guarded_import(name: str, *args, **kwargs):  # noqa: ANN001, ANN202
+    root = name.split(".")[0]
+    if root in BLOCKED_IMPORTS:
+        raise BlockedImport(f"import of {root!r} is not allowed in a solution")
+    return __import__(name, *args, **kwargs)
+
+
+def build_fixtures(spark: SparkSession, fixtures: list[Fixture]) -> dict[str, DataFrame]:
+    """Materialise every fixture as a DataFrame keyed by its declared name."""
+    frames: dict[str, DataFrame] = {}
+    for fx in fixtures:
+        rows = [tuple(r) for r in fx.rows]
+        frames[fx.name] = spark.createDataFrame(rows, schema=fx.schema)
+    return frames
+
+
+def _extract_solve(code: str) -> Any:
+    """exec `code` and hand back its `solve` callable.
+
+    The builtins copy is shallow but private to this call, so swapping
+    __import__ here cannot leak into the host process.
+    """
+    safe_builtins = dict(vars(builtins))
+    safe_builtins["__import__"] = _guarded_import
+    namespace: dict[str, Any] = {"__builtins__": safe_builtins, "__name__": "candidate"}
+
+    exec(compile(code, "<candidate>", "exec"), namespace)  # noqa: S102
+
+    solve = namespace.get("solve")
+    if solve is None or not callable(solve):
+        raise NameError("code does not define a callable named 'solve'")
+    return solve
+
+
+# --------------------------------------------------------------------------
+# Comparison
+# --------------------------------------------------------------------------
+
+
+def _normalise(value: Any, tol: float) -> Any:
+    """Make a collected value comparable and hashable.
+
+    Floats are snapped to a tolerance grid so that two mathematically equal
+    results computed in different partition orders land on the same key.
+    Rows/structs, arrays and maps are flattened recursively.
+    """
+    if value is None:
+        return None
+    if isinstance(value, bool):
+        return value
+    if isinstance(value, float):
+        if math.isnan(value):
+            return "__nan__"
+        if math.isinf(value):
+            return f"__inf_{'pos' if value > 0 else 'neg'}__"
+        if tol > 0:
+            return round(value / tol) * tol
+        return value
+    if isinstance(value, (list, tuple)):
+        return tuple(_normalise(v, tol) for v in value)
+    if isinstance(value, dict):
+        return tuple(sorted((k, _normalise(v, tol)) for k, v in value.items()))
+    if hasattr(value, "asDict"):  # pyspark Row
+        return tuple(
+            sorted((k, _normalise(v, tol)) for k, v in value.asDict(recursive=True).items())
+        )
+    return value
+
+
+def _sort_key(row: tuple) -> tuple:
+    """Total order over heterogeneous rows, including None.
+
+    Sorting by the raw value blows up the moment a column mixes None with an
+    int, which is exactly the case null-handling tasks are built around. Keying
+    on (type name, repr) is stable and never raises.
+    """
+    return tuple((v is None, type(v).__name__, repr(v)) for v in row)
+
+
+def _schema_signature(df: DataFrame, check_order: bool) -> Any:
+    pairs = [(f.name, f.dataType.simpleString()) for f in df.schema.fields]
+    return pairs if check_order else sorted(pairs)
+
+
+def compare_frames(
+    expected: DataFrame, actual: DataFrame, cmp: Compare
+) -> ExecResult:
+    """Judge a candidate DataFrame against the reference DataFrame."""
+    if cmp.check_schema:
+        exp_sig = _schema_signature(expected, cmp.check_column_order)
+        act_sig = _schema_signature(actual, cmp.check_column_order)
+        if exp_sig != act_sig:
+            return ExecResult(
+                False,
+                "schema_mismatch",
+                f"expected {exp_sig}, got {act_sig}",
+            )
+
+    # Align column order before collecting so that a correct answer that simply
+    # selected columns in a different order is not scored as wrong rows.
+    if not cmp.check_column_order and set(expected.columns) == set(actual.columns):
+        actual = actual.select(*expected.columns)
+
+    tol = cmp.float_tolerance
+    exp_rows = [tuple(_normalise(v, tol) for v in r) for r in expected.collect()]
+    act_rows = [tuple(_normalise(v, tol) for v in r) for r in actual.collect()]
+
+    if len(exp_rows) != len(act_rows):
+        return ExecResult(
+            False,
+            "row_mismatch",
+            f"expected {len(exp_rows)} rows, got {len(act_rows)}",
+        )
+
+    if cmp.mode == "rows":
+        exp_rows = sorted(exp_rows, key=_sort_key)
+        act_rows = sorted(act_rows, key=_sort_key)
+
+    for i, (e, a) in enumerate(zip(exp_rows, act_rows, strict=True)):
+        if e != a:
+            return ExecResult(
+                False,
+                "row_mismatch",
+                f"first difference at row {i}: expected {e!r}, got {a!r}",
+            )
+
+    return ExecResult(True, "pass")
+
+
+# --------------------------------------------------------------------------
+# Entry points
+# --------------------------------------------------------------------------
+
+
+def run_code(
+    spark: SparkSession, task: Task, code: str, timeout: int = 60
+) -> tuple[ExecResult, DataFrame | None]:
+    """Run one candidate against one task. Never raises on candidate errors."""
+    try:
+        with _time_limit(timeout):
+            solve = _extract_solve(code)
+            frames = build_fixtures(spark, task.fixtures)
+            result = solve(spark, **frames)
+            if not isinstance(result, DataFrame):
+                return (
+                    ExecResult(
+                        False,
+                        "wrong_type",
+                        f"solve() returned {type(result).__name__}, expected DataFrame",
+                    ),
+                    None,
+                )
+            # Force evaluation inside the time limit: Spark is lazy, so a
+            # candidate that builds a broken plan would otherwise "pass" here
+            # and explode later during comparison.
+            result.cache()
+            result.count()
+            return ExecResult(True, "pass"), result
+    except TaskTimeout as exc:
+        return ExecResult(False, "timeout", str(exc)), None
+    except BlockedImport as exc:
+        return ExecResult(False, "blocked_import", str(exc)), None
+    except NameError as exc:
+        if "solve" in str(exc):
+            return ExecResult(False, "no_solve", str(exc)), None
+        return ExecResult(False, "error", f"{type(exc).__name__}: {exc}"), None
+    except Exception as exc:  # noqa: BLE001 - candidate code, anything goes
+        detail = str(exc).strip().splitlines()
+        head = detail[0] if detail else ""
+        return ExecResult(False, "error", f"{type(exc).__name__}: {head[:400]}"), None
+
+
+def evaluate_candidate(
+    spark: SparkSession, task: Task, code: str, timeout: int = 60
+) -> ExecResult:
+    """Full pipeline for one candidate: run reference, run candidate, compare."""
+    ref_result, expected = run_code(spark, task, task.solution, timeout)
+    if not ref_result.ok or expected is None:
+        # This is a bug in the benchmark, not in the model. Surface it loudly.
+        return ExecResult(
+            False,
+            "reference_broken",
+            f"task {task.id}: reference solution failed: {ref_result.detail}",
+        )
+
+    cand_result, actual = run_code(spark, task, code, timeout)
+    if not cand_result.ok or actual is None:
+        return cand_result
+
+    return compare_frames(expected, actual, task.compare)

join_anti_null_key.yaml

@@ -0,0 +1,39 @@
+id: join_anti_null_key
+category: joins
+difficulty: hard
+probes: >
+  NULL is not equal to NULL in a join predicate. A left-anti join therefore
+  KEEPS left rows whose join key is NULL, because they never matched. Models
+  routinely reason about anti-join as "set difference" and drop them.
+tags: [anti_join, null_semantics]
+
+prompt: |
+  From `orders`, return every order that has no matching customer in `customers`,
+  joined on customer_id. Return all columns of `orders` unchanged.
+
+fixtures:
+  - name: orders
+    schema: order_id INT, customer_id INT, amount DOUBLE
+    rows:
+      - [1, 100, 50.0]
+      - [2, 101, 75.5]
+      - [3, 999, 20.0]
+      - [4, null, 10.0]
+      - [5, 100, 5.25]
+      - [6, null, 42.0]
+  - name: customers
+    schema: customer_id INT, name STRING
+    rows:
+      - [100, "ada"]
+      - [101, "grace"]
+      - [102, "alan"]
+
+solution: |
+  def solve(spark, orders, customers):
+      # left_anti keeps left rows with no match. NULL customer_id never matches
+      # anything (NULL = NULL is unknown, not true), so rows 4 and 6 survive.
+      return orders.join(customers, on="customer_id", how="left_anti")
+
+compare:
+  mode: rows
+  check_schema: false

join_fanout_duplicate_keys.yaml

@@ -0,0 +1,44 @@
+id: join_fanout_duplicate_keys
+category: joins
+difficulty: medium
+probes: >
+  An inner join on a non-unique key multiplies rows. The correct total here is
+  a sum over the fanned-out rows, not over the original left rows. Models that
+  mentally model the join as a lookup produce the un-multiplied number.
+tags: [inner_join, fanout, aggregation]
+
+prompt: |
+  Join `sales` to `rates` on region, then return one row per region with columns
+  region and total, where total is the sum of amount * multiplier across every
+  matched pair. Regions in `sales` with no match in `rates` must be excluded.
+
+fixtures:
+  - name: sales
+    schema: sale_id INT, region STRING, amount DOUBLE
+    rows:
+      - [1, "east", 100.0]
+      - [2, "east", 200.0]
+      - [3, "west", 50.0]
+      - [4, "north", 10.0]
+  - name: rates
+    schema: region STRING, multiplier DOUBLE
+    rows:
+      - ["east", 1.0]
+      - ["east", 2.0]
+      - ["west", 3.0]
+
+solution: |
+  from pyspark.sql import functions as F
+
+  def solve(spark, sales, rates):
+      # east has 2 sales x 2 rates = 4 pairs; the sum must cover all of them.
+      joined = sales.join(rates, on="region", how="inner")
+      return (
+          joined
+          .groupBy("region")
+          .agg(F.sum(F.col("amount") * F.col("multiplier")).alias("total"))
+      )
+
+compare:
+  mode: rows
+  float_tolerance: 1.0e-6

models.py

@@ -0,0 +1,150 @@
+"""Model adapters.
+
+Three backends cover everything you actually need to score:
+
+  ollama:<tag>       local Ollama server (your own models, GGUF quants)
+  openai:<model>     any OpenAI-compatible /v1/chat/completions endpoint,
+                     which includes vLLM, llama.cpp server, TGI, OpenRouter,
+                     and the hosted frontier APIs -- set OPENAI_BASE_URL
+  dummy:<mode>       no inference; for testing the harness itself
+
+Adding a backend means implementing one method. Keep it that way.
+"""
+
+from __future__ import annotations
+
+import json
+import os
+import urllib.error
+import urllib.request
+from abc import ABC, abstractmethod
+
+from .prompting import SYSTEM_PROMPT
+
+
+class ModelError(RuntimeError):
+    pass
+
+
+class Model(ABC):
+    name: str
+
+    @abstractmethod
+    def generate(self, prompt: str, temperature: float, max_tokens: int) -> str:
+        """Return the raw text response. Adapters do not extract code."""
+
+
+def _post_json(url: str, payload: dict, headers: dict, timeout: int) -> dict:
+    req = urllib.request.Request(
+        url,
+        data=json.dumps(payload).encode(),
+        headers={"Content-Type": "application/json", **headers},
+        method="POST",
+    )
+    try:
+        with urllib.request.urlopen(req, timeout=timeout) as resp:
+            return json.loads(resp.read().decode())
+    except urllib.error.HTTPError as exc:
+        body = exc.read().decode(errors="replace")[:500]
+        raise ModelError(f"HTTP {exc.code} from {url}: {body}") from exc
+    except urllib.error.URLError as exc:
+        raise ModelError(f"cannot reach {url}: {exc.reason}") from exc
+
+
+class OllamaModel(Model):
+    def __init__(self, tag: str, host: str | None = None, timeout: int = 300):
+        self.name = f"ollama:{tag}"
+        self.tag = tag
+        self.host = (host or os.environ.get("OLLAMA_HOST") or "http://localhost:11434").rstrip("/")
+        if not self.host.startswith("http"):
+            self.host = f"http://{self.host}"
+        self.timeout = timeout
+
+    def generate(self, prompt: str, temperature: float, max_tokens: int) -> str:
+        payload = {
+            "model": self.tag,
+            "messages": [
+                {"role": "system", "content": SYSTEM_PROMPT},
+                {"role": "user", "content": prompt},
+            ],
+            "stream": False,
+            "options": {
+                "temperature": temperature,
+                "num_predict": max_tokens,
+                # Long fixtures plus a reasoning preamble overflow the 2k
+                # default and the model silently loses the task statement.
+                "num_ctx": 8192,
+            },
+        }
+        data = _post_json(f"{self.host}/api/chat", payload, {}, self.timeout)
+        return data.get("message", {}).get("content", "")
+
+
+class OpenAICompatModel(Model):
+    def __init__(self, model: str, base_url: str | None = None, timeout: int = 300):
+        self.name = f"openai:{model}"
+        self.model = model
+        self.base_url = (
+            base_url or os.environ.get("OPENAI_BASE_URL") or "https://api.openai.com/v1"
+        ).rstrip("/")
+        self.api_key = os.environ.get("OPENAI_API_KEY", "")
+        self.timeout = timeout
+
+    def generate(self, prompt: str, temperature: float, max_tokens: int) -> str:
+        payload = {
+            "model": self.model,
+            "messages": [
+                {"role": "system", "content": SYSTEM_PROMPT},
+                {"role": "user", "content": prompt},
+            ],
+            "temperature": temperature,
+            "max_tokens": max_tokens,
+        }
+        headers = {"Authorization": f"Bearer {self.api_key}"} if self.api_key else {}
+        data = _post_json(
+            f"{self.base_url}/chat/completions", payload, headers, self.timeout
+        )
+        try:
+            return data["choices"][0]["message"]["content"] or ""
+        except (KeyError, IndexError) as exc:
+            raise ModelError(f"unexpected response shape: {str(data)[:300]}") from exc
+
+
+class DummyModel(Model):
+    """Harness self-tests. Never talks to a model.
+
+    reference  -> echo the gold solution. Every task must pass; if one fails,
+                  the benchmark itself is broken.
+    empty      -> return nothing. Every task must fail.
+    """
+
+    def __init__(self, mode: str = "reference"):
+        self.name = f"dummy:{mode}"
+        self.mode = mode
+        self._solutions: dict[str, str] = {}
+
+    def register(self, prompt_key: str, solution: str) -> None:
+        self._solutions[prompt_key] = solution
+
+    def generate(self, prompt: str, temperature: float, max_tokens: int) -> str:
+        if self.mode == "reference":
+            return f"```python\n{self._solutions.get(prompt, '')}\n```"
+        return ""
+
+
+def build_model(spec: str, timeout: int = 300) -> Model:
+    """Parse a `backend:name` spec into a Model."""
+    if ":" not in spec:
+        raise ValueError(
+            f"model spec {spec!r} must look like 'ollama:qwen3:4b' or 'openai:gpt-4o-mini'"
+        )
+    backend, _, rest = spec.partition(":")
+    backend = backend.lower()
+
+    if backend == "ollama":
+        return OllamaModel(rest, timeout=timeout)
+    if backend in ("openai", "vllm", "openai-compat"):
+        return OpenAICompatModel(rest, timeout=timeout)
+    if backend == "dummy":
+        return DummyModel(rest or "reference")
+    raise ValueError(f"unknown backend {backend!r}")

nested_explode_outer_empty.yaml

@@ -0,0 +1,33 @@
+id: nested_explode_outer_empty
+category: schema_nested
+difficulty: medium
+probes: >
+  explode() drops rows whose array is empty or null; explode_outer() keeps them
+  with a NULL element. The fixture contains one empty array and one null array,
+  so the two functions give different row counts.
+tags: [explode, arrays, null_semantics]
+
+prompt: |
+  Flatten the tags array in `docs` so there is one row per tag, KEEPING documents
+  that have an empty or null tags array (their tag should be null).
+
+  Return columns: doc_id, tag.
+
+fixtures:
+  - name: docs
+    schema: doc_id INT, tags ARRAY<STRING>
+    rows:
+      - [1, ["x", "y"]]
+      - [2, []]
+      - [3, null]
+      - [4, ["z"]]
+
+solution: |
+  from pyspark.sql import functions as F
+
+  def solve(spark, docs):
+      # explode() would silently drop docs 2 and 3.
+      return docs.select("doc_id", F.explode_outer("tags").alias("tag"))
+
+compare:
+  mode: rows

null_safe_equality_join.yaml

@@ -0,0 +1,41 @@
+id: null_safe_equality_join
+category: nulls_types
+difficulty: hard
+probes: >
+  Matching NULL to NULL requires the null-safe equality operator (<=> /
+  eqNullSafe). A plain == silently drops the NULL-keyed pair. This is the
+  inverse of join_anti_null_key: same operator, opposite required behaviour,
+  so a model cannot pattern-match its way through both.
+tags: [null_semantics, eqNullSafe, join]
+
+prompt: |
+  Join `left_t` to `right_t` on the code column, treating NULL as a value that
+  matches NULL. Return columns: code, lval, rval, for matched pairs only.
+
+fixtures:
+  - name: left_t
+    schema: code STRING, lval INT
+    rows:
+      - ["a", 1]
+      - ["b", 2]
+      - [null, 3]
+  - name: right_t
+    schema: code STRING, rval INT
+    rows:
+      - ["a", 10]
+      - ["c", 20]
+      - [null, 30]
+
+solution: |
+  from pyspark.sql import functions as F
+
+  def solve(spark, left_t, right_t):
+      # eqNullSafe: NULL <=> NULL is true, so the (null, 3)/(null, 30) pair joins.
+      cond = left_t["code"].eqNullSafe(right_t["code"])
+      return (
+          left_t.join(right_t, cond, "inner")
+          .select(left_t["code"].alias("code"), "lval", "rval")
+      )
+
+compare:
+  mode: rows

null_sum_all_null_group.yaml

@@ -0,0 +1,42 @@
+id: null_sum_all_null_group
+category: nulls_types
+difficulty: medium
+probes: >
+  sum() over a group where every value is NULL returns NULL, not 0. The fix is
+  coalesce after the aggregate, not before -- coalescing before changes the
+  average and the count. Models very often emit sum(coalesce(x,0)) which is a
+  different query that happens to agree here but not on avg.
+tags: [null_semantics, sum, coalesce]
+
+prompt: |
+  For each region in `readings`, return the sum of `value` as `total`, but
+  report 0 rather than null when the region has no non-null values at all.
+  `total` must be a LONG (BIGINT) column.
+
+  Return columns: region, total.
+
+fixtures:
+  - name: readings
+    schema: region STRING, value INT
+    rows:
+      - ["north", 5]
+      - ["north", null]
+      - ["north", 7]
+      - ["south", null]
+      - ["south", null]
+      - ["east", 3]
+
+solution: |
+  from pyspark.sql import functions as F
+
+  def solve(spark, readings):
+      # sum() ignores NULLs; an all-NULL group aggregates to NULL, so the
+      # coalesce has to sit outside the aggregate.
+      return (
+          readings
+          .groupBy("region")
+          .agg(F.coalesce(F.sum("value"), F.lit(0)).cast("long").alias("total"))
+      )
+
+compare:
+  mode: rows

prompting.py

@@ -0,0 +1,92 @@
+"""Turning a task into a model prompt, and model output back into code.
+
+Extraction is deliberately forgiving. A model that writes correct PySpark but
+wraps it in prose should score as correct -- otherwise the benchmark is partly
+measuring instruction-following formatting, which is a different axis and one
+that would flatter models tuned on this exact style.
+"""
+
+from __future__ import annotations
+
+import re
+
+from .schema import Task
+
+SYSTEM_PROMPT = (
+    "You are an expert PySpark engineer. You write correct, idiomatic PySpark "
+    "using the DataFrame API. You respond with a single Python code block and "
+    "no explanation."
+)
+
+_TEMPLATE = """{prompt}
+
+Input DataFrames (already created, passed as arguments):
+{fixtures}
+
+Write a single Python function with exactly this signature:
+
+    def solve(spark, {args}):
+        ...
+        return result
+
+Requirements:
+- Return a PySpark DataFrame.
+- Put any imports you need inside the code block (e.g. `from pyspark.sql import functions as F`).
+- Do not create your own data. Use only the DataFrames passed in.
+- Do not call .show(), .collect(), or print().
+"""
+
+
+def describe_fixtures(task: Task) -> str:
+    lines = []
+    for fx in task.fixtures:
+        lines.append(f"  {fx.name}: {fx.schema}")
+    return "\n".join(lines)
+
+
+def build_prompt(task: Task) -> str:
+    """The user-turn text for a task."""
+    return _TEMPLATE.format(
+        prompt=task.prompt,
+        fixtures=describe_fixtures(task),
+        args=", ".join(task.fixture_names),
+    )
+
+
+_FENCE_RE = re.compile(
+    r"```(?:python|py)?\s*\n(.*?)(?:```|\Z)", re.DOTALL | re.IGNORECASE
+)
+_THINK_RE = re.compile(r"<think>.*?</think>", re.DOTALL | re.IGNORECASE)
+
+
+def extract_code(text: str) -> str:
+    """Pull runnable Python out of a model response.
+
+    Handles, in order: reasoning-model <think> blocks, fenced code blocks
+    (picking the one that actually defines solve), and bare code.
+    """
+    if not text:
+        return ""
+
+    # Reasoning traces routinely contain draft code that does not run. Strip
+    # them before looking for the answer, or we score the scratchpad.
+    text = _THINK_RE.sub("", text)
+    # An unterminated <think> means the model ran out of budget mid-reasoning.
+    if "<think>" in text.lower():
+        text = re.sub(r"<think>.*\Z", "", text, flags=re.DOTALL | re.IGNORECASE)
+
+    blocks = [b.strip() for b in _FENCE_RE.findall(text)]
+    if blocks:
+        for block in blocks:
+            if "def solve(" in block:
+                return block
+        return blocks[0]
+
+    if "def solve(" in text:
+        # Bare code, no fence. Drop any prose before the first import/def.
+        lines = text.splitlines()
+        for i, line in enumerate(lines):
+            if re.match(r"^\s*(from|import|def)\s", line):
+                return "\n".join(lines[i:]).strip()
+
+    return text.strip()

runner.py

@@ -0,0 +1,302 @@
+"""Run a model over the suite and score it.
+
+pass@k uses the unbiased estimator from Chen et al. (2021), "Evaluating Large
+Language Models Trained on Code" -- not the naive "did any of k samples pass",
+which is biased upward and not comparable across different n.
+"""
+
+from __future__ import annotations
+
+import json
+import time
+from collections import defaultdict
+from dataclasses import asdict, dataclass, field
+from datetime import datetime, timezone
+from pathlib import Path
+
+from pyspark.sql import SparkSession
+
+from .harness import ExecResult, evaluate_candidate
+from .models import DummyModel, Model, ModelError
+from .prompting import build_prompt, extract_code
+from .schema import Task
+
+
+def pass_at_k(n: int, c: int, k: int) -> float:
+    """Probability that at least one of k samples drawn from n passes.
+
+    n = samples generated, c = samples that passed.
+    """
+    if n < k:
+        raise ValueError(f"cannot estimate pass@{k} from only {n} samples")
+    if n - c < k:
+        return 1.0
+    # Product form avoids overflow in the binomial coefficients.
+    prob = 1.0
+    for i in range(k):
+        prob *= (n - c - i) / (n - i)
+    return 1.0 - prob
+
+
+@dataclass
+class SampleRecord:
+    task_id: str
+    category: str
+    difficulty: str
+    sample_index: int
+    passed: bool
+    status: str
+    detail: str
+    raw_response: str
+    extracted_code: str
+    latency_s: float
+
+
+@dataclass
+class TaskRecord:
+    task_id: str
+    category: str
+    difficulty: str
+    probes: str
+    n: int
+    c: int
+    statuses: dict[str, int] = field(default_factory=dict)
+
+
+@dataclass
+class RunReport:
+    model: str
+    n_samples: int
+    temperature: float
+    started_at: str
+    duration_s: float
+    n_tasks: int
+    pass_at_1: float
+    pass_at_k: dict[str, float]
+    by_category: dict[str, dict]
+    by_difficulty: dict[str, dict]
+    failure_modes: dict[str, int]
+    tasks: list[TaskRecord]
+    samples: list[SampleRecord]
+
+    def to_json(self) -> str:
+        return json.dumps(asdict(self), indent=2)
+
+
+def get_spark(app_name: str = "spark-eval") -> SparkSession:
+    """A small, deterministic, local Spark session.
+
+    Single shuffle partition is deliberate: it makes float aggregation order
+    reproducible and cuts per-task overhead by more than half. Tasks are tiny;
+    parallelism buys nothing here.
+    """
+    return (
+        SparkSession.builder.appName(app_name)
+        .master("local[2]")
+        .config("spark.sql.shuffle.partitions", "1")
+        .config("spark.default.parallelism", "2")
+        .config("spark.sql.adaptive.enabled", "false")
+        .config("spark.ui.enabled", "false")
+        .config("spark.sql.session.timeZone", "UTC")
+        .config("spark.driver.memory", "2g")
+        .getOrCreate()
+    )
+
+
+def run_suite(
+    tasks: list[Task],
+    model: Model,
+    *,
+    n_samples: int = 1,
+    ks: tuple[int, ...] = (1,),
+    temperature: float = 0.2,
+    max_tokens: int = 1024,
+    timeout: int = 60,
+    spark: SparkSession | None = None,
+    keep_responses: bool = True,
+    progress: bool = True,
+) -> RunReport:
+    owns_spark = spark is None
+    spark = spark or get_spark()
+    spark.sparkContext.setLogLevel("ERROR")
+
+    started = datetime.now(timezone.utc)
+    t0 = time.time()
+
+    samples: list[SampleRecord] = []
+    task_records: list[TaskRecord] = []
+    failure_modes: dict[str, int] = defaultdict(int)
+
+    for idx, task in enumerate(tasks, 1):
+        prompt = build_prompt(task)
+        # The dummy backend needs the gold answer keyed by the exact prompt.
+        if isinstance(model, DummyModel):
+            model.register(prompt, task.solution)
+
+        statuses: dict[str, int] = defaultdict(int)
+        passed_count = 0
+
+        for s in range(n_samples):
+            s_t0 = time.time()
+            try:
+                raw = model.generate(prompt, temperature, max_tokens)
+            except ModelError as exc:
+                raw = ""
+                result = ExecResult(False, "model_error", str(exc))
+                code = ""
+            else:
+                code = extract_code(raw)
+                if not code.strip():
+                    result = ExecResult(False, "empty_response", "no code in response")
+                else:
+                    result = evaluate_candidate(spark, task, code, timeout=timeout)
+
+            latency = time.time() - s_t0
+            statuses[result.status] += 1
+            if result.ok:
+                passed_count += 1
+            else:
+                failure_modes[result.status] += 1
+
+            samples.append(
+                SampleRecord(
+                    task_id=task.id,
+                    category=task.category,
+                    difficulty=task.difficulty,
+                    sample_index=s,
+                    passed=result.ok,
+                    status=result.status,
+                    detail=result.detail,
+                    raw_response=raw if keep_responses else "",
+                    extracted_code=code if keep_responses else "",
+                    latency_s=round(latency, 3),
+                )
+            )
+
+            # A broken reference solution means the benchmark is lying. Stop.
+            if result.status == "reference_broken":
+                raise RuntimeError(result.detail)
+
+        task_records.append(
+            TaskRecord(
+                task_id=task.id,
+                category=task.category,
+                difficulty=task.difficulty,
+                probes=task.probes,
+                n=n_samples,
+                c=passed_count,
+                statuses=dict(statuses),
+            )
+        )
+
+        if progress:
+            mark = "PASS" if passed_count == n_samples else (
+                "FAIL" if passed_count == 0 else f"{passed_count}/{n_samples}"
+            )
+            print(
+                f"[{idx:>3}/{len(tasks)}] {task.id:<40} {mark}",
+                flush=True,
+            )
+
+    def _agg(records: list[TaskRecord]) -> dict:
+        if not records:
+            return {"n_tasks": 0, "pass_at_1": 0.0}
+        out = {
+            "n_tasks": len(records),
+            "pass_at_1": round(
+                sum(pass_at_k(r.n, r.c, 1) for r in records) / len(records), 4
+            ),
+        }
+        for k in ks:
+            if k > 1 and n_samples >= k:
+                out[f"pass_at_{k}"] = round(
+                    sum(pass_at_k(r.n, r.c, k) for r in records) / len(records), 4
+                )
+        return out
+
+    by_category: dict[str, dict] = {}
+    for cat in sorted({r.category for r in task_records}):
+        by_category[cat] = _agg([r for r in task_records if r.category == cat])
+
+    by_difficulty: dict[str, dict] = {}
+    for diff in ("easy", "medium", "hard"):
+        subset = [r for r in task_records if r.difficulty == diff]
+        if subset:
+            by_difficulty[diff] = _agg(subset)
+
+    overall = _agg(task_records)
+    report = RunReport(
+        model=model.name,
+        n_samples=n_samples,
+        temperature=temperature,
+        started_at=started.isoformat(),
+        duration_s=round(time.time() - t0, 2),
+        n_tasks=len(task_records),
+        pass_at_1=overall["pass_at_1"],
+        pass_at_k={
+            f"pass_at_{k}": overall[f"pass_at_{k}"]
+            for k in ks
+            if k > 1 and f"pass_at_{k}" in overall
+        },
+        by_category=by_category,
+        by_difficulty=by_difficulty,
+        failure_modes=dict(sorted(failure_modes.items(), key=lambda kv: -kv[1])),
+        tasks=task_records,
+        samples=samples,
+    )
+
+    if owns_spark:
+        spark.stop()
+
+    return report
+
+
+def format_report(report: RunReport) -> str:
+    """Human-readable summary. The JSON is the machine-readable artifact."""
+    lines = [
+        "",
+        "=" * 68,
+        f"  spark-eval  |  {report.model}",
+        "=" * 68,
+        f"  tasks       {report.n_tasks}",
+        f"  samples     {report.n_samples} per task @ temperature {report.temperature}",
+        f"  duration    {report.duration_s}s",
+        "",
+        f"  pass@1      {report.pass_at_1:.1%}",
+    ]
+    for key, v in report.pass_at_k.items():
+        label = "pass@" + key.rsplit("_", 1)[-1]
+        lines.append(f"  {label:<11} {v:.1%}")
+
+    def _plural(n: int) -> str:
+        return f"{n} task" if n == 1 else f"{n} tasks"
+
+    # Worst category first: the point of the breakdown is finding the weakness.
+    lines += ["", "  By category", "  " + "-" * 46]
+    for cat, stats in sorted(
+        report.by_category.items(), key=lambda kv: kv[1]["pass_at_1"]
+    ):
+        lines.append(
+            f"    {cat:<22} {stats['pass_at_1']:>6.1%}  ({_plural(stats['n_tasks'])})"
+        )
+
+    if report.by_difficulty:
+        lines += ["", "  By difficulty", "  " + "-" * 46]
+        for diff, stats in report.by_difficulty.items():
+            lines.append(
+                f"    {diff:<22} {stats['pass_at_1']:>6.1%}  ({_plural(stats['n_tasks'])})"
+            )
+
+    if report.failure_modes:
+        lines += ["", "  Failure modes", "  " + "-" * 46]
+        total = sum(report.failure_modes.values())
+        for mode, count in report.failure_modes.items():
+            lines.append(f"    {mode:<22} {count:>4}  ({count / total:.0%})")
+
+    lines += ["", "=" * 68, ""]
+    return "\n".join(lines)
+
+
+def write_report(report: RunReport, path: Path) -> None:
+    path.parent.mkdir(parents=True, exist_ok=True)
+    path.write_text(report.to_json())

schema.py

@@ -0,0 +1,247 @@
+"""Task schema, loading, and validation.
+
+A task is a YAML file. The contract is deliberately narrow so that a task is
+cheap to write by hand and impossible to score ambiguously:
+
+  - `fixtures` declare the input DataFrames by schema + literal rows. They are
+    small, deterministic, and committed to the repo. No network, no generated
+    data, no randomness.
+  - `prompt` is what the model sees. It names the fixtures and states the
+    required entrypoint signature.
+  - `solution` is reference PySpark that a human wrote and that the harness
+    executes to produce the expected output. There is no hardcoded expected
+    table anywhere -- expected output is *computed*, so a fixture edit can
+    never silently desynchronise from a stale golden file.
+  - `compare` says how to judge equality. Default is order-insensitive rows
+    plus exact schema.
+
+Every task must define `solve(spark, **frames) -> DataFrame`.
+"""
+
+from __future__ import annotations
+
+import re
+from dataclasses import dataclass, field
+from pathlib import Path
+from typing import Any, Literal
+
+import yaml
+
+CATEGORIES = {
+    "joins",
+    "windows",
+    "aggregations",
+    "schema_nested",
+    "udf_vs_native",
+    "nulls_types",
+    "sql_translation",
+    "delta_merge",
+}
+
+DIFFICULTIES = {"easy", "medium", "hard"}
+
+_ID_RE = re.compile(r"^[a-z0-9]+(?:_[a-z0-9]+)*$")
+
+
+class TaskValidationError(ValueError):
+    """Raised when a task file is structurally invalid."""
+
+
+@dataclass(frozen=True)
+class Fixture:
+    """One input DataFrame, defined literally.
+
+    `schema` is a Spark DDL string (e.g. "id INT, name STRING"). We use DDL
+    rather than inferring from rows because inference silently changes types
+    when a column happens to be all-null in the sample, and null handling is
+    one of the things this benchmark is trying to measure.
+    """
+
+    name: str
+    schema: str
+    rows: list[list[Any]]
+
+    def __post_init__(self) -> None:
+        if not self.name.isidentifier():
+            raise TaskValidationError(
+                f"fixture name {self.name!r} is not a valid Python identifier"
+            )
+        if not self.schema.strip():
+            raise TaskValidationError(f"fixture {self.name!r} has an empty schema")
+
+
+@dataclass(frozen=True)
+class Compare:
+    """How to decide whether a candidate result matches the reference."""
+
+    # "rows" -> order-insensitive multiset comparison (the default; most tasks
+    #           do not specify an order, so requiring one would fail correct code)
+    # "ordered_rows" -> order matters (use when the prompt explicitly asks for
+    #           a sort, e.g. window/top-n tasks)
+    mode: Literal["rows", "ordered_rows"] = "rows"
+
+    # Exact schema match (names, types, nullability-insensitive). Turning this
+    # off is a deliberate loosening -- record why in the task file.
+    check_schema: bool = True
+
+    # Column names must match exactly and in order. Off means we compare on the
+    # set of columns, useful when the prompt does not pin an output column order.
+    check_column_order: bool = False
+
+    # Absolute tolerance for float/double columns. Spark's floating point
+    # aggregation order is not deterministic across partitions, so exact
+    # equality on doubles is a flaky-test generator.
+    float_tolerance: float = 1e-9
+
+
+@dataclass(frozen=True)
+class Task:
+    id: str
+    category: str
+    difficulty: str
+    prompt: str
+    fixtures: list[Fixture]
+    solution: str
+    compare: Compare = field(default_factory=Compare)
+    # Free-text note on what this task is actually probing. Shows up in the
+    # per-category failure report; the point of the benchmark is diagnosis,
+    # not just a number.
+    probes: str = ""
+    tags: list[str] = field(default_factory=list)
+    source_path: Path | None = None
+
+    @property
+    def fixture_names(self) -> list[str]:
+        return [f.name for f in self.fixtures]
+
+
+def _require(data: dict, key: str, path: Path, type_: type) -> Any:
+    if key not in data:
+        raise TaskValidationError(f"{path}: missing required key {key!r}")
+    value = data[key]
+    if not isinstance(value, type_):
+        raise TaskValidationError(
+            f"{path}: key {key!r} must be {type_.__name__}, got {type(value).__name__}"
+        )
+    return value
+
+
+def load_task(path: Path) -> Task:
+    """Parse and validate a single task file."""
+    with path.open() as fh:
+        raw = yaml.safe_load(fh)
+
+    if not isinstance(raw, dict):
+        raise TaskValidationError(f"{path}: top level must be a mapping")
+
+    task_id = _require(raw, "id", path, str)
+    if not _ID_RE.match(task_id):
+        raise TaskValidationError(
+            f"{path}: id {task_id!r} must be lower_snake_case"
+        )
+
+    category = _require(raw, "category", path, str)
+    if category not in CATEGORIES:
+        raise TaskValidationError(
+            f"{path}: unknown category {category!r}; expected one of {sorted(CATEGORIES)}"
+        )
+
+    difficulty = raw.get("difficulty", "medium")
+    if difficulty not in DIFFICULTIES:
+        raise TaskValidationError(
+            f"{path}: difficulty {difficulty!r} must be one of {sorted(DIFFICULTIES)}"
+        )
+
+    prompt = _require(raw, "prompt", path, str).strip()
+    solution = _require(raw, "solution", path, str)
+
+    raw_fixtures = _require(raw, "fixtures", path, list)
+    if not raw_fixtures:
+        raise TaskValidationError(f"{path}: at least one fixture is required")
+
+    fixtures = []
+    for item in raw_fixtures:
+        if not isinstance(item, dict):
+            raise TaskValidationError(f"{path}: each fixture must be a mapping")
+        fixtures.append(
+            Fixture(
+                name=_require(item, "name", path, str),
+                schema=_require(item, "schema", path, str),
+                rows=[list(r) for r in _require(item, "rows", path, list)],
+            )
+        )
+
+    names = [f.name for f in fixtures]
+    if len(set(names)) != len(names):
+        raise TaskValidationError(f"{path}: duplicate fixture names in {names}")
+
+    raw_compare = raw.get("compare") or {}
+    if not isinstance(raw_compare, dict):
+        raise TaskValidationError(f"{path}: 'compare' must be a mapping")
+    unknown = set(raw_compare) - {
+        "mode",
+        "check_schema",
+        "check_column_order",
+        "float_tolerance",
+    }
+    if unknown:
+        raise TaskValidationError(f"{path}: unknown compare keys {sorted(unknown)}")
+    compare = Compare(**raw_compare)
+    if compare.mode not in ("rows", "ordered_rows"):
+        raise TaskValidationError(f"{path}: invalid compare.mode {compare.mode!r}")
+
+    # The reference solution has to honour the same contract we ask of models.
+    if "def solve(" not in solution:
+        raise TaskValidationError(
+            f"{path}: solution must define solve(spark, ...); "
+            "the harness calls it by name"
+        )
+
+    # A prompt that does not mention a fixture is a prompt the model cannot
+    # answer. This has caught more authoring bugs than any other check.
+    for name in names:
+        if name not in prompt:
+            raise TaskValidationError(
+                f"{path}: fixture {name!r} is never mentioned in the prompt"
+            )
+
+    return Task(
+        id=task_id,
+        category=category,
+        difficulty=difficulty,
+        prompt=prompt,
+        fixtures=fixtures,
+        solution=solution,
+        compare=compare,
+        probes=raw.get("probes", ""),
+        tags=list(raw.get("tags", [])),
+        source_path=path,
+    )
+
+
+def load_tasks(
+    root: Path,
+    categories: list[str] | None = None,
+    ids: list[str] | None = None,
+) -> list[Task]:
+    """Load every task under `root`, optionally filtered.
+
+    Sorted by id so that runs are reproducible and diffable.
+    """
+    paths = sorted(root.rglob("*.yaml")) + sorted(root.rglob("*.yml"))
+    tasks = [load_task(p) for p in paths]
+
+    seen: dict[str, Path] = {}
+    for t in tasks:
+        if t.id in seen:
+            raise TaskValidationError(
+                f"duplicate task id {t.id!r} in {t.source_path} and {seen[t.id]}"
+            )
+        seen[t.id] = t.source_path  # type: ignore[assignment]
+
+    if categories:
+        tasks = [t for t in tasks if t.category in set(categories)]
+    if ids:
+        tasks = [t for t in tasks if t.id in set(ids)]
+
+    return sorted(tasks, key=lambda t: t.id)

sql_having_to_dataframe.yaml

@@ -0,0 +1,54 @@
+id: sql_having_to_dataframe
+category: sql_translation
+difficulty: medium
+probes: >
+  HAVING filters after aggregation, WHERE filters before. The query has both,
+  with different predicates, so swapping them changes the answer. Tests whether
+  the model understands the clause order rather than pattern-matching filter().
+tags: [sql, having, where, groupby]
+
+prompt: |
+  Translate this SQL into the PySpark DataFrame API against `txns`
+  (do not use spark.sql or createOrReplaceTempView):
+
+      SELECT store, SUM(amount) AS total
+      FROM txns
+      WHERE status = 'ok'
+      GROUP BY store
+      HAVING COUNT(*) >= 2
+      ORDER BY total DESC
+
+  Return columns: store, total.
+
+fixtures:
+  - name: txns
+    schema: txn_id INT, store STRING, status STRING, amount DOUBLE
+    rows:
+      - [1, "s1", "ok", 100.0]
+      - [2, "s1", "ok", 50.0]
+      - [3, "s1", "void", 999.0]
+      - [4, "s2", "ok", 400.0]
+      - [5, "s2", "void", 1.0]
+      - [6, "s3", "ok", 10.0]
+      - [7, "s3", "ok", 20.0]
+      - [8, "s3", "ok", 30.0]
+
+solution: |
+  from pyspark.sql import functions as F
+
+  def solve(spark, txns):
+      # WHERE before groupBy, HAVING as a filter on the aggregated frame.
+      # s2 has only one 'ok' row, so HAVING COUNT(*) >= 2 removes it.
+      return (
+          txns
+          .filter(F.col("status") == "ok")
+          .groupBy("store")
+          .agg(F.sum("amount").alias("total"), F.count(F.lit(1)).alias("_n"))
+          .filter(F.col("_n") >= 2)
+          .drop("_n")
+          .orderBy(F.col("total").desc())
+      )
+
+compare:
+  mode: ordered_rows
+  float_tolerance: 1.0e-6

test_harness.py

@@ -0,0 +1,202 @@
+"""Unit tests for the comparator and the execution guards.
+
+The comparator is the thing everything else trusts, so its edge cases get
+tested directly rather than only through tasks.
+"""
+
+from __future__ import annotations
+
+import pytest
+
+from spark_eval.harness import compare_frames, run_code
+from spark_eval.runner import pass_at_k
+from spark_eval.schema import Compare, Fixture, Task
+
+
+def _frame(spark, schema, rows):
+    return spark.createDataFrame([tuple(r) for r in rows], schema=schema)
+
+
+def _task(solution: str, compare: Compare | None = None) -> Task:
+    return Task(
+        id="t",
+        category="joins",
+        difficulty="easy",
+        prompt="uses df",
+        fixtures=[Fixture(name="df", schema="a INT", rows=[[1], [2]])],
+        solution=solution,
+        compare=compare or Compare(),
+    )
+
+
+# --------------------------------------------------------------------------
+# Comparison semantics
+# --------------------------------------------------------------------------
+
+
+def test_row_order_ignored_by_default(spark):
+    a = _frame(spark, "x INT", [[1], [2], [3]])
+    b = _frame(spark, "x INT", [[3], [1], [2]])
+    assert compare_frames(a, b, Compare(mode="rows")).ok
+
+
+def test_row_order_enforced_when_requested(spark):
+    a = _frame(spark, "x INT", [[1], [2], [3]])
+    b = _frame(spark, "x INT", [[3], [1], [2]])
+    assert not compare_frames(a, b, Compare(mode="ordered_rows")).ok
+
+
+def test_nulls_sort_without_raising(spark):
+    """A column mixing NULL and INT used to blow up the sort key."""
+    a = _frame(spark, "x INT", [[1], [None], [3]])
+    b = _frame(spark, "x INT", [[None], [3], [1]])
+    assert compare_frames(a, b, Compare(mode="rows")).ok
+
+
+def test_null_is_not_equal_to_zero(spark):
+    a = _frame(spark, "x INT", [[None]])
+    b = _frame(spark, "x INT", [[0]])
+    assert not compare_frames(a, b, Compare()).ok
+
+
+def test_schema_mismatch_detected(spark):
+    a = _frame(spark, "x INT", [[1]])
+    b = _frame(spark, "x BIGINT", [[1]])
+    result = compare_frames(a, b, Compare(check_schema=True))
+    assert not result.ok
+    assert result.status == "schema_mismatch"
+
+
+def test_schema_check_can_be_relaxed(spark):
+    a = _frame(spark, "x INT", [[1]])
+    b = _frame(spark, "x BIGINT", [[1]])
+    assert compare_frames(a, b, Compare(check_schema=False)).ok
+
+
+def test_column_order_ignored_by_default(spark):
+    a = _frame(spark, "x INT, y INT", [[1, 2]])
+    b = _frame(spark, "y INT, x INT", [[2, 1]])
+    assert compare_frames(a, b, Compare()).ok
+
+
+def test_column_order_enforced_when_requested(spark):
+    a = _frame(spark, "x INT, y INT", [[1, 2]])
+    b = _frame(spark, "y INT, x INT", [[2, 1]])
+    assert not compare_frames(a, b, Compare(check_column_order=True)).ok
+
+
+def test_float_tolerance_absorbs_partition_order(spark):
+    a = _frame(spark, "x DOUBLE", [[0.1 + 0.2]])
+    b = _frame(spark, "x DOUBLE", [[0.3]])
+    assert compare_frames(a, b, Compare(float_tolerance=1e-6)).ok
+    assert not compare_frames(a, b, Compare(float_tolerance=0.0)).ok
+
+
+def test_row_count_mismatch_reported(spark):
+    a = _frame(spark, "x INT", [[1], [2]])
+    b = _frame(spark, "x INT", [[1]])
+    result = compare_frames(a, b, Compare())
+    assert result.status == "row_mismatch"
+    assert "expected 2 rows, got 1" in result.detail
+
+
+def test_duplicate_rows_are_significant(spark):
+    """Multiset, not set: a dropDuplicates bug must not pass."""
+    a = _frame(spark, "x INT", [[1], [1], [2]])
+    b = _frame(spark, "x INT", [[1], [2], [2]])
+    assert not compare_frames(a, b, Compare()).ok
+
+
+def test_nested_struct_compared_recursively(spark):
+    schema = "s STRUCT<a: INT, b: STRING>"
+    a = _frame(spark, schema, [[(1, "x")]])
+    b = _frame(spark, schema, [[(1, "y")]])
+    assert not compare_frames(a, b, Compare()).ok
+    assert compare_frames(a, a, Compare()).ok
+
+
+def test_array_order_is_significant(spark):
+    a = _frame(spark, "xs ARRAY<INT>", [[[1, 2]]])
+    b = _frame(spark, "xs ARRAY<INT>", [[[2, 1]]])
+    assert not compare_frames(a, b, Compare()).ok
+
+
+# --------------------------------------------------------------------------
+# Execution guards
+# --------------------------------------------------------------------------
+
+
+def test_missing_solve_reported_cleanly(spark):
+    result, _ = run_code(spark, _task("def other(): pass"), "def other(): pass")
+    assert result.status == "no_solve"
+
+
+def test_non_dataframe_return_reported(spark):
+    code = "def solve(spark, df):\n    return 42\n"
+    result, _ = run_code(spark, _task(code), code)
+    assert result.status == "wrong_type"
+
+
+def test_syntax_error_is_a_failure_not_a_crash(spark):
+    code = "def solve(spark, df:\n    return df\n"
+    result, _ = run_code(spark, _task(code), code)
+    assert result.status == "error"
+    assert "SyntaxError" in result.detail
+
+
+def test_blocked_import_is_rejected(spark):
+    code = "import os\n\ndef solve(spark, df):\n    return df\n"
+    result, _ = run_code(spark, _task(code), code)
+    assert result.status == "blocked_import"
+
+
+def test_lazy_failure_is_caught_at_run_time(spark):
+    """Spark is lazy; a broken plan must fail here, not later in compare."""
+    code = (
+        "from pyspark.sql import functions as F\n"
+        "def solve(spark, df):\n"
+        "    return df.select(F.col('does_not_exist'))\n"
+    )
+    result, _ = run_code(spark, _task(code), code)
+    assert result.status == "error"
+
+
+def test_timeout_is_enforced(spark):
+    code = (
+        "def solve(spark, df):\n"
+        "    while True:\n"
+        "        pass\n"
+    )
+    result, _ = run_code(spark, _task(code), code, timeout=2)
+    assert result.status == "timeout"
+
+
+# --------------------------------------------------------------------------
+# pass@k estimator
+# --------------------------------------------------------------------------
+
+
+@pytest.mark.parametrize(
+    ("n", "c", "k", "expected"),
+    [
+        (1, 1, 1, 1.0),
+        (1, 0, 1, 0.0),
+        (10, 0, 1, 0.0),
+        (10, 10, 5, 1.0),
+        (10, 5, 1, 0.5),
+        # 2 of 4 pass; drawing 2 of 4 misses both only 1 time in 6.
+        (4, 2, 2, 1 - (2 / 4) * (1 / 3)),
+    ],
+)
+def test_pass_at_k_values(n, c, k, expected):
+    assert pass_at_k(n, c, k) == pytest.approx(expected)
+
+
+def test_pass_at_k_rejects_k_greater_than_n():
+    with pytest.raises(ValueError, match="cannot estimate"):
+        pass_at_k(2, 1, 5)
+
+
+def test_pass_at_k_is_monotonic_in_k():
+    values = [pass_at_k(10, 3, k) for k in range(1, 8)]
+    assert values == sorted(values)

test_mutants.py

@@ -0,0 +1,188 @@
+"""Mutation tests: prove the harness fails plausible-but-wrong code.
+
+A benchmark that only checks "the gold solution passes" is worthless -- an
+always-return-True comparator would satisfy that. Every task here gets a
+*mutant*: the specific wrong implementation a model actually tends to produce.
+The suite asserts the harness rejects it.
+
+If you add a task, add its mutant. CI enforces the pairing.
+"""
+
+from __future__ import annotations
+
+import pytest
+
+from spark_eval.harness import evaluate_candidate
+from spark_eval.schema import load_tasks
+
+from .conftest import TASKS_DIR
+
+# task_id -> (description of the mistake, wrong implementation)
+MUTANTS: dict[str, tuple[str, str]] = {
+    "join_anti_null_key": (
+        "treats anti-join as set difference, dropping NULL-keyed rows",
+        """
+from pyspark.sql import functions as F
+
+def solve(spark, orders, customers):
+    ids = [r[0] for r in customers.select("customer_id").collect()]
+    return orders.filter(~F.col("customer_id").isin(ids))
+""",
+    ),
+    "join_fanout_duplicate_keys": (
+        "dedups the right side first, so the fanout never happens",
+        """
+from pyspark.sql import functions as F
+
+def solve(spark, sales, rates):
+    r = rates.dropDuplicates(["region"])
+    return (sales.join(r, on="region", how="inner")
+                 .groupBy("region")
+                 .agg(F.sum(F.col("amount") * F.col("multiplier")).alias("total")))
+""",
+    ),
+    "window_default_frame_ties": (
+        "uses an explicit ROWS frame instead of the default RANGE frame",
+        """
+from pyspark.sql import functions as F
+from pyspark.sql.window import Window
+
+def solve(spark, events):
+    w = Window.partitionBy("user").orderBy("ts").rowsBetween(
+        Window.unboundedPreceding, Window.currentRow
+    )
+    return events.withColumn("running", F.sum("value").over(w))
+""",
+    ),
+    "window_rank_family_ties": (
+        "uses row_number() for all three columns, erasing tie behaviour",
+        """
+from pyspark.sql import functions as F
+from pyspark.sql.window import Window
+
+def solve(spark, scores):
+    w = Window.partitionBy("team").orderBy(F.col("points").desc(), F.col("player").asc())
+    return (scores
+            .withColumn("rnk", F.row_number().over(w))
+            .withColumn("dense", F.row_number().over(w))
+            .withColumn("rownum", F.row_number().over(w))
+            .orderBy(F.col("team").asc(), F.col("points").desc(), F.col("player").asc()))
+""",
+    ),
+    "agg_count_null_semantics": (
+        "uses count('*') everywhere, ignoring NULL and DISTINCT semantics",
+        """
+from pyspark.sql import functions as F
+
+def solve(spark, staff):
+    return staff.groupBy("dept").agg(
+        F.count(F.lit(1)).alias("n_rows"),
+        F.count(F.lit(1)).alias("n_emails"),
+        F.count(F.lit(1)).alias("n_distinct"),
+    )
+""",
+    ),
+    "null_sum_all_null_group": (
+        "plain sum(), so the all-NULL group returns NULL instead of 0",
+        """
+from pyspark.sql import functions as F
+
+def solve(spark, readings):
+    return readings.groupBy("region").agg(F.sum("value").cast("long").alias("total"))
+""",
+    ),
+    "null_safe_equality_join": (
+        "plain equality, which silently drops the NULL/NULL pair",
+        """
+from pyspark.sql import functions as F
+
+def solve(spark, left_t, right_t):
+    return (left_t.join(right_t, left_t["code"] == right_t["code"], "inner")
+                  .select(left_t["code"].alias("code"), "lval", "rval"))
+""",
+    ),
+    "nested_explode_outer_empty": (
+        "explode() instead of explode_outer(), dropping empty/null arrays",
+        """
+from pyspark.sql import functions as F
+
+def solve(spark, docs):
+    return docs.select("doc_id", F.explode("tags").alias("tag"))
+""",
+    ),
+    "upsert_latest_version": (
+        "overwrite instead of upsert: target-only rows are lost",
+        """
+from pyspark.sql import functions as F
+from pyspark.sql.window import Window
+
+def solve(spark, target, updates):
+    w = Window.partitionBy("id").orderBy(F.col("version").desc())
+    return (updates.withColumn("_rn", F.row_number().over(w))
+                   .filter(F.col("_rn") == 1)
+                   .drop("_rn"))
+""",
+    ),
+    "udf_null_input_handling": (
+        "no None guard, so the UDF raises on the NULL row",
+        """
+from pyspark.sql import functions as F
+from pyspark.sql.types import IntegerType
+
+def solve(spark, notes):
+    word_udf = F.udf(lambda s: len(s.split()), IntegerType())
+    return notes.withColumn("n_words", word_udf(F.col("text")))
+""",
+    ),
+    "sql_having_to_dataframe": (
+        "applies HAVING as a WHERE, filtering rows instead of groups",
+        """
+from pyspark.sql import functions as F
+
+def solve(spark, txns):
+    return (txns
+            .filter(F.col("status") == "ok")
+            .groupBy("store")
+            .agg(F.sum("amount").alias("total"))
+            .orderBy(F.col("total").desc()))
+""",
+    ),
+    "agg_pivot_fill": (
+        "forgets fillna, leaving NULL for absent store/quarter combinations",
+        """
+from pyspark.sql import functions as F
+
+def solve(spark, sales_long):
+    return (sales_long.groupBy("store")
+                      .pivot("quarter", ["Q1", "Q2", "Q3"])
+                      .agg(F.sum("amount")))
+""",
+    ),
+}
+
+ALL_TASKS = {t.id: t for t in load_tasks(TASKS_DIR)}
+
+
+def test_every_task_has_a_mutant():
+    """Keeps the two files honest with each other."""
+    missing = sorted(set(ALL_TASKS) - set(MUTANTS))
+    orphaned = sorted(set(MUTANTS) - set(ALL_TASKS))
+    assert not missing, f"tasks with no mutant test: {missing}"
+    assert not orphaned, f"mutants for tasks that no longer exist: {orphaned}"
+
+
+@pytest.mark.parametrize("task_id", sorted(MUTANTS))
+def test_mutant_is_rejected(spark, task_id):
+    task = ALL_TASKS[task_id]
+    description, code = MUTANTS[task_id]
+    result = evaluate_candidate(spark, task, code, timeout=90)
+
+    assert result.status != "reference_broken", (
+        f"{task_id}: the reference solution itself failed -- fix the task, "
+        f"not the mutant ({result.detail})"
+    )
+    assert not result.ok, (
+        f"{task_id}: harness ACCEPTED a wrong answer ({description}). "
+        f"The task cannot distinguish correct from incorrect code; "
+        f"strengthen the fixtures."
+    )

test_prompting.py

@@ -0,0 +1,81 @@
+"""Extraction tests.
+
+Every case here is a real response shape seen from a local model. Extraction
+bugs show up as a uniform score drop that looks like a model being bad, which
+is the most expensive kind of benchmark bug to debug.
+"""
+
+from __future__ import annotations
+
+from spark_eval.prompting import build_prompt, extract_code
+from spark_eval.schema import Fixture, Task
+
+SOLVE = "def solve(spark, df):\n    return df"
+
+
+def test_fenced_python_block():
+    assert extract_code(f"Here you go:\n\n```python\n{SOLVE}\n```\n") == SOLVE
+
+
+def test_fence_without_language_tag():
+    assert extract_code(f"```\n{SOLVE}\n```") == SOLVE
+
+
+def test_unterminated_fence_still_extracts():
+    """Hitting the token limit mid-block should not cost the model the task."""
+    assert extract_code(f"```python\n{SOLVE}") == SOLVE
+
+
+def test_picks_the_block_that_defines_solve():
+    text = (
+        "First, the schema:\n\n```python\nschema = 'a INT'\n```\n\n"
+        f"And the answer:\n\n```python\n{SOLVE}\n```"
+    )
+    assert extract_code(text) == SOLVE
+
+
+def test_reasoning_block_is_stripped():
+    text = (
+        "<think>\nMaybe ```python\ndef solve(spark, df): return None\n```\n"
+        "no wait, that is wrong.\n</think>\n\n"
+        f"```python\n{SOLVE}\n```"
+    )
+    assert extract_code(text) == SOLVE
+
+
+def test_unterminated_reasoning_block_yields_nothing_runnable():
+    """Ran out of budget while thinking -- must not score the scratchpad."""
+    code = extract_code("<think>\nI should probably write\n```python\nx = 1\n```")
+    assert "def solve" not in code
+
+
+def test_bare_code_without_fence():
+    text = f"Sure, this works.\n{SOLVE}"
+    assert extract_code(text) == SOLVE
+
+
+def test_empty_response():
+    assert extract_code("") == ""
+    assert extract_code("   \n  ") == ""
+
+
+def test_prose_only_response_has_no_solve():
+    assert "def solve" not in extract_code("I cannot help with that request.")
+
+
+def test_prompt_names_every_fixture_and_the_signature():
+    task = Task(
+        id="t",
+        category="joins",
+        difficulty="easy",
+        prompt="Join orders to customers.",
+        fixtures=[
+            Fixture(name="orders", schema="id INT", rows=[[1]]),
+            Fixture(name="customers", schema="id INT", rows=[[1]]),
+        ],
+        solution=SOLVE,
+    )
+    prompt = build_prompt(task)
+    assert "def solve(spark, orders, customers)" in prompt
+    assert "orders: id INT" in prompt
+    assert "customers: id INT" in prompt

udf_null_input_handling.yaml

@@ -0,0 +1,42 @@
+id: udf_null_input_handling
+category: udf_vs_native
+difficulty: medium
+probes: >
+  A Python UDF is called with None for NULL input rather than being skipped, so
+  an unguarded body raises and fails the whole stage. Also pins the return type:
+  an unannotated UDF defaults to StringType and silently stringifies integers.
+tags: [udf, null_semantics, return_type]
+
+prompt: |
+  Add a column `n_words` to `notes` containing the number of whitespace-separated
+  words in `text`. Rows where text is null must get null (not 0, not an error).
+  `n_words` must be an INT column.
+
+  Return columns: note_id, text, n_words.
+
+fixtures:
+  - name: notes
+    schema: note_id INT, text STRING
+    rows:
+      - [1, "hello world"]
+      - [2, null]
+      - [3, "one"]
+      - [4, "  padded   words here  "]
+      - [5, ""]
+
+solution: |
+  from pyspark.sql import functions as F
+  from pyspark.sql.types import IntegerType
+
+  def solve(spark, notes):
+      # The None guard is mandatory: Spark hands NULL to the UDF as None.
+      def count_words(s):
+          if s is None:
+              return None
+          return len(s.split())
+
+      word_udf = F.udf(count_words, IntegerType())
+      return notes.withColumn("n_words", word_udf(F.col("text")))
+
+compare:
+  mode: rows

upsert_latest_version.yaml

@@ -0,0 +1,56 @@
+id: upsert_latest_version
+category: delta_merge
+difficulty: medium
+probes: >
+  The core of a MERGE/upsert, expressed without the Delta dependency: collapse
+  a change stream to the latest row per key, then union the keys that only
+  exist in the target. The trap is dropping target rows that the source never
+  touched, which is the difference between an upsert and an overwrite.
+tags: [upsert, scd, dedup, window]
+
+prompt: |
+  `target` is a current-state table and `updates` is a change stream that may
+  contain several rows per id. Produce the post-upsert state:
+
+    - for an id present in `updates`, keep only the row with the highest version
+      from `updates` (it wins over `target` regardless of target's version)
+    - for an id absent from `updates`, keep the `target` row unchanged
+
+  Return columns: id, value, version.
+
+fixtures:
+  - name: target
+    schema: id INT, value STRING, version INT
+    rows:
+      - [1, "old_a", 1]
+      - [2, "old_b", 5]
+      - [3, "old_c", 2]
+  - name: updates
+    schema: id INT, value STRING, version INT
+    rows:
+      - [1, "new_a1", 2]
+      - [1, "new_a2", 3]
+      - [2, "new_b", 4]
+      - [4, "new_d", 1]
+
+solution: |
+  from pyspark.sql import functions as F
+  from pyspark.sql.window import Window
+
+  def solve(spark, target, updates):
+      # 1. collapse the change stream to one winning row per id
+      w = Window.partitionBy("id").orderBy(F.col("version").desc())
+      latest = (
+          updates
+          .withColumn("_rn", F.row_number().over(w))
+          .filter(F.col("_rn") == 1)
+          .drop("_rn")
+      )
+      # 2. keep target rows the stream never mentions. Note id=2 takes the
+      #    update (version 4) even though target's version 5 is higher: the
+      #    spec says the source wins, which is what MERGE ... WHEN MATCHED does.
+      untouched = target.join(latest.select("id"), on="id", how="left_anti")
+      return untouched.unionByName(latest)
+
+compare:
+  mode: rows

window_default_frame_ties.yaml

@@ -0,0 +1,42 @@
+id: window_default_frame_ties
+category: windows
+difficulty: hard
+probes: >
+  The single most-missed detail in Spark windows. With an ORDER BY and no
+  explicit frame, the default is RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT
+  ROW -- a *value* range, so tied ordering keys all see each other's rows.
+  Writing rowsBetween(Window.unboundedPreceding, 0) gives a different, wrong
+  answer on ties. Both look identical in a code review.
+tags: [window, frame, range_vs_rows, ties]
+
+prompt: |
+  For each row in `events`, compute a running total of value within each user,
+  ordered by ts, using the SQL default window frame (RANGE BETWEEN UNBOUNDED
+  PRECEDING AND CURRENT ROW). Rows that share the same ts within a user must
+  therefore share the same running total.
+
+  Return columns: user, ts, value, running.
+
+fixtures:
+  - name: events
+    schema: user STRING, ts INT, value INT
+    rows:
+      - ["a", 1, 10]
+      - ["a", 2, 20]
+      - ["a", 2, 30]
+      - ["a", 3, 40]
+      - ["b", 1, 5]
+      - ["b", 1, 7]
+
+solution: |
+  from pyspark.sql import functions as F
+  from pyspark.sql.window import Window
+
+  def solve(spark, events):
+      # No rowsBetween/rangeBetween call: this is the SQL default frame,
+      # which is RANGE-based. The two ts=2 rows for user 'a' both get 60.
+      w = Window.partitionBy("user").orderBy("ts")
+      return events.withColumn("running", F.sum("value").over(w))
+
+compare:
+  mode: rows

window_rank_family_ties.yaml

@@ -0,0 +1,51 @@
+id: window_rank_family_ties
+category: windows
+difficulty: medium
+probes: >
+  rank / dense_rank / row_number diverge only in the presence of ties, and
+  rank leaves gaps while dense_rank does not. Fixtures are built so that all
+  three produce different columns.
+tags: [window, rank, dense_rank, row_number]
+
+prompt: |
+  For each row in `scores`, add three columns computed over a window
+  partitioned by team and ordered by points descending:
+    - rnk         using rank()
+    - dense       using dense_rank()
+    - rownum      using row_number()
+
+  Return columns: team, player, points, rnk, dense, rownum.
+  Order the output by team ascending, then points descending, then player ascending.
+
+fixtures:
+  - name: scores
+    schema: team STRING, player STRING, points INT
+    rows:
+      - ["red", "ann", 10]
+      - ["red", "bob", 10]
+      - ["red", "cal", 7]
+      - ["red", "dee", 5]
+      - ["blue", "eve", 3]
+      - ["blue", "fay", 3]
+
+solution: |
+  from pyspark.sql import functions as F
+  from pyspark.sql.window import Window
+
+  def solve(spark, scores):
+      w = Window.partitionBy("team").orderBy(F.col("points").desc())
+      # row_number needs a deterministic tiebreak or the result is not stable
+      # across runs; the prompt pins the output order, so break on player.
+      w_rn = Window.partitionBy("team").orderBy(
+          F.col("points").desc(), F.col("player").asc()
+      )
+      return (
+          scores
+          .withColumn("rnk", F.rank().over(w))
+          .withColumn("dense", F.dense_rank().over(w))
+          .withColumn("rownum", F.row_number().over(w_rn))
+          .orderBy(F.col("team").asc(), F.col("points").desc(), F.col("player").asc())
+      )
+
+compare:
+  mode: ordered_rows